package leetcode.tree.binary.remove;

import leetcode.tree.binary.TreeNode;

/**
 * 814. 二叉树剪枝
 *
 * 移除所有不包含 1的子树 (全是0的子树)
 */
public class PruneTree {

    public TreeNode pruneTree(TreeNode root) {
        if (isZeroTree(root)) {
            return null;
        } else {
            return root;
        }
    }

    /**
     * 递归实现 自己写的 真开心 打败了100%的用户
     *
     * 方法作用
     *      判断当前树的节点是否都是 0
     *      条件 = 左子树和右子树都是 0并且当前根节点也是 0
     *
     * @param root 以 root为根的子树
     * @return  当前树是否全为 0
     */
    public boolean isZeroTree(TreeNode root) {
        // 叶子节点的左节点或者右节点
        if (root == null) {
            return true;
        }
        // 叶子节点
        if (root.left == null && root.right == null) {
            return root.val == 0;
        }

        // 判断左子树是否都为0
        boolean left = isZeroTree(root.left);
        // 是的话 删除左子树
        if (left) {
            root.left = null;
        }
        boolean right = isZeroTree(root.right);
        if (right) {
            root.right = null;
        }

        // 左右子树和当前节点的值都为0 这棵树才是都是0
        return left && right && root.val == 0;
    }
}
